package com.mybatisRule;

import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.scripting.xmltags.SqlNode;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.util.*;

/**
 * SQL 校验left * 等语法
 */
@Component
public class MyBatisRuleApplicationRunner implements ApplicationRunner {

    @Autowired
    private SqlSessionFactory sqlSessionFactory;

    @Value("${mybatisRule.checkFlag:false}")
    private boolean checkFlag;

    @Value("${mybatisRule.whiteList:}")
    private List<String> whiteList;

    @Override
    public void run(ApplicationArguments args) throws Exception {
        if (!checkFlag) {
            System.out.println("####################################################");
            System.out.println("######           跳过SQL规范检查!!!!!!         ######");
            System.out.println("####################################################");

            return;
        }
        System.out.println("####################################################");
        System.out.println("####              SQL规范检查开始             ########");
        System.out.println("####################################################");

        int errorCount = 0;
        Configuration configuration = sqlSessionFactory.getConfiguration();
        Collection<MappedStatement> mappedStatements = configuration.getMappedStatements();
        try {
            for (Object m : mappedStatements) {
                MappedStatement mm = null;
                if (m instanceof MappedStatement) {
                    mm = (MappedStatement) m;
                } else {
                    continue;
                }
                String id = mm.getId();
                if (!SqlCommandType.SELECT.equals(mm.getSqlCommandType())) {
                    continue;
                }
                SqlSource sqlSource = mm.getSqlSource();
                if (sqlSource == null) {
                    continue;
                }
                Class<?> sqlSourceClass = sqlSource.getClass();
                Optional<Field> rootSqlNodeAny = Arrays.stream(sqlSourceClass.getDeclaredFields()).filter(field -> field.getName().equals("rootSqlNode")).findAny();
                if (rootSqlNodeAny.isPresent()) {
                    Field rootSqlNodeField = rootSqlNodeAny.get();
                    if (rootSqlNodeField == null) {
                        continue;
                    }
                    rootSqlNodeField.setAccessible(true);
                    Object rootSqlNode = rootSqlNodeField.get(sqlSource);
                    if (rootSqlNode == null) {
                        continue;
                    }
                    Class<?> rootSqlNodeClass = rootSqlNode.getClass();
                    Optional<Field> contentsAny = Arrays.stream(rootSqlNodeClass.getDeclaredFields()).filter(field -> field.getName().equals("contents")).findAny();
                    if (contentsAny.isPresent()) {
                        Field contentsField = contentsAny.get();
                        contentsField.setAccessible(true);
                        List<SqlNode> contents = (List<SqlNode>) contentsField.get(rootSqlNode);
                        StringBuilder sb = new StringBuilder();
                        for (SqlNode sqlNode : contents) {
                            Class<?> sqlNodeClass = sqlNode.getClass();
                            Optional<Field> textAny = Arrays.stream(sqlNodeClass.getDeclaredFields()).filter(field -> field.getName().equals("text")).findAny();
                            Optional<Field> sqlAny = Arrays.stream(sqlNodeClass.getDeclaredFields()).filter(field -> field.getName().equals("sql")).findAny();
                            if (textAny.isPresent()) {
                                Field textField = textAny.get();
                                textField.setAccessible(true);
                                Object o = textField.get(sqlNode);
                                sb.append(o);
                            } else if (sqlAny.isPresent()) {
                                Field sqlField = sqlAny.get();
                                sqlField.setAccessible(true);
                                Object o = sqlField.get(sqlNode);
                                sb.append(o);
                            }
                        }

                        if (sb.toString().toLowerCase(Locale.ROOT).contains(" join ") || sb.toString().contains("*")) {
                            if (whiteList.contains(id)) {
                                System.err.println("跳过SQLID: " + id);
                                continue;
                            }
                            System.err.println(++errorCount + "> SQLID: " + id + sb.toString());
                        }
                    } else {
                        continue;
                    }
                } else {
                    continue;
                }
            }
        } catch (Exception e) {
            System.err.println(e);
        }

        if (errorCount > 0) {
            System.err.println("#############################################################");
            System.err.println("######     以上SQL语法不规范!!! 禁止join,子查询,*号等操作      ####");
            System.err.println("######                    程序退出中....                 ######");
            System.err.println("#############################################################");
            System.exit(0);
        } else {
            System.err.println("#############################################################");
            System.err.println("##########             SQL规范检查通过            ##############");
            System.err.println("#############################################################");
        }
    }
}

